#include <ntndk.h>

typedef struct 
{ 
	ULONG     Length;
	ULONG     Unknown[11];
} RTL_HEAP_DEFINITION, *PRTL_HEAP_DEFINITION;
BOOL IsPEFile(HANDLE g_hHeap, HANDLE hfile)
{
	IO_STATUS_BLOCK objIoStatus;  
	NTSTATUS	ntStatus; 
	IMAGE_DOS_HEADER * pdoshead;
	IMAGE_NT_HEADERS * pnthead;
	unsigned char * pBuffer;
	BOOL RETDATA;
	pBuffer = (unsigned char *)RtlAllocateHeap( g_hHeap, 0, 1024 );
	NtReadFile(hfile,0,0,0,&objIoStatus,pBuffer,1024,0,0);
	pdoshead = pBuffer;
	if (pdoshead->e_magic != 0x5a4d)
	{
		RETDATA= FALSE;
		goto IsPEFileend;
	}
	pnthead = pdoshead->e_lfanew+pBuffer;
	if (pnthead->Signature != 0x4550)
	{
		RETDATA= FALSE;
		goto IsPEFileend;
	}
	RETDATA = TRUE;
IsPEFileend:
	RtlFreeHeap( g_hHeap, 0, pBuffer );
	return RETDATA;
}

BOOL ProcessFile(HANDLE g_hHeap, PUNICODE_STRING pfilename)

{
	HANDLE hfile,hHeap;
	UNICODE_STRING  filename;
	OBJECT_ATTRIBUTES obAttr;
	IO_STATUS_BLOCK objIoStatus;  
	NTSTATUS	ntStatus; 
	RTL_HEAP_DEFINITION  heapParams;
//	FILE_BOTH_DIR_INFORMATION * pBuffer;
	unsigned char * pBuffer;
	obAttr.Length = sizeof(obAttr);
	obAttr.Attributes = OBJ_CASE_INSENSITIVE;
	obAttr.ObjectName = pfilename;
	obAttr.RootDirectory = 0 ;
	obAttr.SecurityDescriptor = 0 ;
	obAttr.SecurityQualityOfService = 0 ;
	pBuffer = (unsigned char *)RtlAllocateHeap( g_hHeap, 0, 1024 );
	ntStatus = NtOpenFile(&hfile, SYNCHRONIZE|FILE_EXECUTE|FILE_READ_DATA,
		&obAttr, &objIoStatus,
		FILE_SHARE_DELETE|FILE_SHARE_READ,
		FILE_SYNCHRONOUS_IO_NONALERT|FILE_NON_DIRECTORY_FILE);

	if (ntStatus==0)
	{
		if (IsPEFile(g_hHeap,hfile))
		{
			return TRUE;
		}
//		NtReadFile(hfile,0,0,0,&objIoStatus,pBuffer,1024,0,0);
		NtClose(hfile);
	}
	RtlFreeHeap( g_hHeap, 0, pBuffer );
	return FALSE;
}
void ProcessDir(HANDLE g_hHeap, PUNICODE_STRING pfilename)
{
	HANDLE hfinder,hHeap;
	UNICODE_STRING  filename;
	OBJECT_ATTRIBUTES obAttr;
	IO_STATUS_BLOCK objIoStatus;  
	NTSTATUS	ntStatus; 
	RTL_HEAP_DEFINITION  heapParams;
	FILE_BOTH_DIR_INFORMATION * pBuffer;
	WCHAR * sznewDirname;
	obAttr.Length = sizeof(obAttr);
	obAttr.Attributes = OBJ_CASE_INSENSITIVE;
	obAttr.ObjectName = pfilename;
	obAttr.RootDirectory = 0 ;
	obAttr.SecurityDescriptor = 0 ;
	obAttr.SecurityQualityOfService = 0 ;
	pBuffer = (FILE_BOTH_DIR_INFORMATION*)RtlAllocateHeap( g_hHeap, 0, 1024 );
	sznewDirname = (WCHAR*)RtlAllocateHeap( g_hHeap, 0, 1024 );
	ntStatus =NtOpenFile(&hfinder,FILE_LIST_DIRECTORY | SYNCHRONIZE | FILE_ANY_ACCESS,
		&obAttr,&objIoStatus,3, 
		FILE_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT | FILE_OPEN_FOR_BACKUP_INTENT);
	if (ntStatus == 0)
	{
		do 
		{
			ntStatus = NtQueryDirectoryFile(hfinder,
				0,
				0,
				0,
				&objIoStatus,
				pBuffer,  
				1024,  
				FileBothDirectoryInformation,  
				1,
				0,
				0) ;
			if (ntStatus==0)
			{
				if ((pBuffer->FileAttributes&FILE_ATTRIBUTE_DIRECTORY)==FILE_ATTRIBUTE_DIRECTORY)
				{
					if (wcscmp(pBuffer->FileName,L"..")&&wcscmp(pBuffer->FileName,L"."))
					{
						memset((unsigned char *)sznewDirname,0,1024);
						wcsncpy(&sznewDirname[wcslen(sznewDirname)],&pfilename->Buffer[0],wcslen(&pfilename->Buffer[0]));
						wcsncpy(&sznewDirname[wcslen(sznewDirname)],L"\\",1);
						wcsncpy(&sznewDirname[wcslen(sznewDirname)],pBuffer->FileName,pBuffer->FileNameLength/2 );
						RtlInitUnicodeString(&filename,sznewDirname);
						ProcessDir(g_hHeap,&filename);
		//				wcsncpy(&sznewDirname[wcslen(sznewDirname)],L"\n",1);
	//					RtlInitUnicodeString(&filename,sznewDirname);
	//					NtDisplayString(&filename);
					}
				}
				else
				{
					memset((unsigned char *)sznewDirname,0,1024);
					wcsncpy(&sznewDirname[wcslen(sznewDirname)],&pfilename->Buffer[0],wcslen(&pfilename->Buffer[0]));
					wcsncpy(&sznewDirname[wcslen(sznewDirname)],L"\\",1);
					wcsncpy(&sznewDirname[wcslen(sznewDirname)],pBuffer->FileName,pBuffer->FileNameLength/2 );
					RtlInitUnicodeString(&filename,sznewDirname);
					if (ProcessFile(g_hHeap,&filename))
					{
						wcsncpy(&sznewDirname[wcslen(sznewDirname)],L"\n",1);
						RtlInitUnicodeString(&filename,sznewDirname);
						NtDisplayString(&filename);

					}

				}
			}
		} while (!ntStatus);

		NtClose(hfinder);
	}
	RtlFreeHeap( g_hHeap, 0, pBuffer );
	RtlFreeHeap( g_hHeap, 0, sznewDirname );

}

void NtProcessStartup(PPEB ppeb)
{
	NTSTATUS status;
//	NtDisplayString()
//	status = ChangeHostName();
	//FILE_BOTH_DIR_INFORMATION
	HANDLE g_hHeap;
	UNICODE_STRING  filename;
	UNICODE_STRING  filename2;
	RTL_HEAP_DEFINITION  heapParams;
	_asm{
//		int 3;
	}
	RtlInitUnicodeString(&filename2,L"OK\n");
	NtDisplayString(&filename2);
	RtlInitUnicodeString(&filename,L"\\??\\C:\\WINDOWS");
	g_hHeap = RtlCreateHeap( 2, 0, 0x100000, 0x1000, 0, &heapParams );
	ProcessDir(g_hHeap,&filename);
	NtTerminateProcess(NtCurrentProcess(),0);
}


